import cdsapi
import os
from datetime import datetime, timedelta


#用于获取训练用数据
# 配置
save_dir = "./era5_data"
os.makedirs(save_dir, exist_ok=True)

# 变量设定
surface_variables = [
    'mean_sea_level_pressure', '10m_u_component_of_wind',
    '10m_v_component_of_wind', '2m_temperature'
]

upper_variables = [
    'geopotential', 'specific_humidity', 'temperature',
    'u_component_of_wind', 'v_component_of_wind'
]

pressure_levels = ['1000', '925', '850', '700', '600', '500',
                   '400', '300', '250', '200', '150', '100', '50']

area = [90, 0, -90, 359.75]  # 全球

# 时间设置
start_date = datetime(2022, 12, 18)
end_date = datetime(2022, 12, 19)  # 可以拓展范围
time_steps = ['{:02d}:00'.format(h) for h in range(24)]  # 每小时

c = cdsapi.Client()

# 下载循环
curr_date = start_date
while curr_date <= end_date:
    date_str = curr_date.strftime("%Y-%m-%d")
    print(f"Downloading {date_str} ...")

    surface_path = os.path.join(save_dir, f"surface_{date_str}.nc")
    upper_path = os.path.join(save_dir, f"upper_{date_str}.nc")

    # Surface level download
    if not os.path.exists(surface_path):
        c.retrieve('reanalysis-era5-single-levels', {
            'product_type': 'reanalysis',
            'format': 'netcdf',
            'variable': surface_variables,
            'date': date_str,
            'time': time_steps,
            'area': area,
        }, surface_path)

    # # Upper air level download
    # if not os.path.exists(upper_path):
    #     c.retrieve('reanalysis-era5-pressure-levels', {
    #         'product_type': 'reanalysis',
    #         'format': 'netcdf',
    #         'variable': upper_variables,
    #         'pressure_level': pressure_levels,
    #         'date': date_str,
    #         'time': time_steps,
    #         'area': area,
    #     }, upper_path)

    curr_date += timedelta(days=1)

print("下载完成")
